import * as tstl from 'typescript-to-lua';
import * as ts from 'typescript';
import * as lua from 'typescript-to-lua/dist/LuaAST';
import { FunctionVisitor, TransformationContext } from 'typescript-to-lua/dist/transformation/context';
import { AnnotationKind, getTypeAnnotations } from 'typescript-to-lua/dist/transformation/utils/annotations';
import { getSymbolExportScope } from "typescript-to-lua/dist/transformation/utils/export";
import { createLocalOrExportedOrGlobalDeclaration } from "typescript-to-lua/dist/transformation/utils/lua-ast";
import { isFirstDeclaration } from "typescript-to-lua/dist/transformation/utils/typescript";
import { transformIdentifier } from "typescript-to-lua/dist/transformation/visitors/identifier";
import { transformPropertyName } from "typescript-to-lua/dist/transformation/visitors/literal";

console.log('Custom enum_rule loaded');

/**
 * 尝试获取常量枚举的编译时常量值
 * @param context 转换上下文
 * @param node 枚举成员节点或属性访问表达式
 * @returns 返回对应的Lua字面量表达式（字符串或数值），如果无法获取常量值则返回 undefined
 */
export function tryGetConstEnumValue(
    context: TransformationContext,
    node: ts.EnumMember | ts.PropertyAccessExpression | ts.ElementAccessExpression
): lua.Expression | undefined {
    // 通过类型检查器获取编译时常量值
    const value = context.checker.getConstantValue(node);
    // 处理字符串类型的常量值
    if (typeof value === "string") {
        return lua.createStringLiteral(value, node);
    } else if (typeof value === "number") {
        // 处理数值类型的常量值
        return lua.createNumericLiteral(value, node);
    }
    // 非编译时常量或无法确定值时返回 undefined
}

const enum_rule = {
    visitors: {
        // 处理枚举成员的情况
        [ts.SyntaxKind.EnumDeclaration]: (node: ts.EnumDeclaration, context) => {
            // 处理常量枚举声明（当不保留常量枚举时直接跳过）
            if (ts.getCombinedModifierFlags(node) & ts.ModifierFlags.Const && !context.options.preserveConstEnums) {
                return undefined;
            }
            const type = context.checker.getTypeAtLocation(node);

            // 检查是否仅需要成员声明（通过 @compileMembersOnly 注解控制）
            const membersOnly = getTypeAnnotations(type).has(AnnotationKind.CompileMembersOnly);
            const result: lua.Statement[] = [];
            const tableEntries = [] as tstl.TableFieldExpression[];

            // 生成枚举表初始结构（仅首次声明且需要完整结构时）
            if (!membersOnly && isFirstDeclaration(context, node)) {
                const name = transformIdentifier(context, node.name);
                // 创建带默认值的 Lua 表：enumTable = enumTable or {}
                const table = lua.createBinaryExpression(
                    lua.cloneIdentifier(name),
                    lua.createTableExpression(),
                    lua.SyntaxKind.OrOperator   
                );
                // 创建表字段 - 使用对象字面量形式
                result.push(...createLocalOrExportedOrGlobalDeclaration(context, name, table, node));
            }
            const enumReference = context.transformExpression(node.name);
            // 遍历处理每个枚举成员
            for (const member of node.members) {
                const memberName = transformPropertyName(context, member.name);
                let valueExpression: lua.Expression | undefined;
                // 尝试获取常量枚举值
                const constEnumValue = tryGetConstEnumValue(context, member);
                if (constEnumValue) {
                    valueExpression = constEnumValue;
                } else if (member.initializer) {
                    // 处理引用其他枚举成员的情况
                    if (ts.isIdentifier(member.initializer)) {
                        const symbol = context.checker.getSymbolAtLocation(member.initializer);
                        if (
                            symbol?.valueDeclaration &&
                            ts.isEnumMember(symbol.valueDeclaration) &&
                            symbol.valueDeclaration.parent === node
                        ) {
                            // 生成其他成员的引用表达式 enum[otherMember]
                            const otherMemberName = transformPropertyName(context, symbol.valueDeclaration.name);
                            valueExpression = lua.createTableIndexExpression(enumReference, otherMemberName);
                        }
                    }
                    // 转换成员初始化表达式
                    if (!valueExpression) {
                        valueExpression = context.transformExpression(member.initializer);
                    }
                } else {
                    // 无初始值时使用 nil 作为默认值
                    valueExpression = lua.createNilLiteral();
                }
                // 根据编译模式生成不同结构
                if (membersOnly) {
                    // 仅生成成员变量声明
                    const enumSymbol = context.checker.getSymbolAtLocation(node.name);
                    const exportScope = enumSymbol ? getSymbolExportScope(context, enumSymbol) : undefined;

                    tableEntries.push({
                        kind: tstl.SyntaxKind.TableFieldExpression as tstl.SyntaxKind,
                        key: tstl.createStringLiteral(memberName) as tstl.StringLiteral,
                        value: valueExpression as tstl.Expression
                    });

                    result.push(
                        ...createLocalOrExportedOrGlobalDeclaration(
                            context,
                            lua.isIdentifier(memberName)
                                ? memberName
                                : lua.createIdentifier(member.name.getText(), member.name),
                            valueExpression,
                            node,
                            exportScope
                        )
                    );
                } else {
                    // 生成完整的枚举表赋值语句 enum.member = value
                    const memberAccessor = lua.createTableIndexExpression(enumReference, memberName);
                    result.push(lua.createAssignmentStatement(memberAccessor, valueExpression, member));
                }
            }
            return result;
        },
    }
};

module.exports = enum_rule;